#!/usr/bin/env python
#
# Copyright (C) 2007 Jason Kivlighn
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, #USA.

import sqlite3
import random
import os, sys

### Options ###
INTERVAL_RANGE = 1 #in weeks
THRESHOLD = 10 #only show artists with at least this many plays
REGRESSION_FACTOR = 2 * INTERVAL_RANGE

WEEK = 7*24*60*60

def a_close_entry(i,list,close=5):
  for j in range(0,close):
    if i+j in list:
      return i+j
    if i-j in list:
      return i-j
  return -1

def raw_data(conn,type,user,interval=INTERVAL_RANGE,threshold=THRESHOLD,stack=True):
  all_data = {}

  #This prevents SQL injection, not to mention that other values are meaningless and going to die
  if type not in ["artists","tracks","albums","track_tags","artist_tags"]:
    type = "artists"
    print "WARNING: Got bad type, defaulting to 'artists'"

  c = conn.cursor()
  
  intervals = []
  c.execute("SELECT MIN(begin) FROM %s" % type)
  begin = c.fetchone()[0]
  c.execute("SELECT MAX(end) FROM %s" % type)
  end = c.fetchone()[0]
  
  c.execute("SELECT DISTINCT(begin) FROM %s" % type)
  non_zero_intervals = [b for b, in c]

  for i in range(begin,end,WEEK):
    closest = a_close_entry(i,non_zero_intervals)
    if closest < 0:
      intervals.append(i)
    else:
      intervals.append(closest)
  
  lengthened_intervals = []
  for i in xrange(0,len(intervals),interval):
    lengthened_intervals.append(intervals[i])
  
  for begin in lengthened_intervals:
    all_data[begin] = []
  
  ignored_artists = []
  c.execute("SELECT %s,SUM(playcount) AS plays FROM %s GROUP BY %s HAVING plays <= %d" % (type[:-1],type,type[:-1],threshold))
  for artist,_ in c:
    ignored_artists.append(artist)
  
  artist_stack = []
  flag = False
  c.execute("SELECT DISTINCT(%s) FROM %s ORDER BY begin" % (type[:-1],type))
  artists = filter(lambda a: a not in ignored_artists,[artist for artist, in c])
  for artist in artists:
    if flag:
      artist_stack.append(artist)
    else:
      artist_stack.insert(0,artist)
    
    flag = not flag
  
  '''
  c.execute("SELECT begin,artist,playcount FROM entries")
  for begin,artist,plays in c:
    all_data[begin].append((artist,int(plays)))
  '''
  
  for x in xrange(len(lengthened_intervals)-1):
    c.execute("SELECT %s,SUM(playcount) FROM %s WHERE begin BETWEEN %d AND %d GROUP BY %s" % (type[:-1],type,lengthened_intervals[x],lengthened_intervals[x+1],type[:-1]))
    for artist,plays in c:
      all_data[lengthened_intervals[x]].append((artist,float(plays)))
  
  data = []
  for begin in sorted(all_data.keys()):
    c.execute("SELECT %s,playcount FROM %s WHERE begin=%d" % (type[:-1],type,begin))
    current_artists = {}
    for a,p in c:
      current_artists[a] = p
    
    curr_stack = [len(data),0]
    
    level = 0
    for artist in artist_stack:
      if artist in current_artists.keys():
        if stack:
          level += current_artists[artist]
        else:
          level = current_artists[artist]
      curr_stack.append(level)
        
    data.append(curr_stack)

  return lengthened_intervals, artist_stack, data

def center_data(data):
  max_playcount = max([max(l[1:]) for l in data])
  center = max_playcount/2
  
  for l in data:
    curr_max = max(l)
    diff = center-curr_max/2
    for i in range(len(l)-1):
      l[i+1] += diff

def linregress_data(conn,type,user,interval=INTERVAL_RANGE,threshold=THRESHOLD, regression_factor=REGRESSION_FACTOR,stack=True,centered=True):
  
  import regression as r
  
  times, artist_stack, data = raw_data(conn,type,user,interval,threshold,stack)
  if centered:
    center_data(data)
  
  #Store all the x-axis values
  curved_data = []
  for x in range(0,len(data)*regression_factor-regression_factor-1):
    curved_data.append([x])
  
  #Store all the y-axis values, obtains through linear regression
  for i in range(len(artist_stack)+1):
    plays = [[x[i+1],x[0]*regression_factor] for x in data]
    coeffs = r.linearRegression(plays,len(plays)-1)
    
    counter = 0
    for x in range(0,len(data)*regression_factor-regression_factor-1):
      y = 0.0
      for c in range(len(coeffs)):
        y += coeffs[c]*pow(x,c)
        
      if y != y:
        curved_data[counter].append(0.0)
      else:
        curved_data[counter].append(y)
        
      counter += 1
  return times, artist_stack, curved_data

def spline_data(conn,type,user,interval=INTERVAL_RANGE,threshold=THRESHOLD,x_interval=0.1,stack=True,centered=True):
  times, artist_stack, data = raw_data(conn,type,user,interval,threshold,stack)
  if centered:
    center_data(data)

  from numpy import r_, array
  from scipy.signal import cspline1d, cspline1d_eval
  
  x = [d[0] for d in data]
  dx = data[0][0] - data[1][0]
  
  spline_data = []
  newys = []
  newx = r_[0:data[-1][0]:x_interval]
  for i in range(len(artist_stack)+1):
    y = array([d[i+1] for d in data])
    cj = cspline1d(y)
    newy = cspline1d_eval(cj, newx, dx=dx,x0=x[0])
    newys.append(newy)
    
  for x in newx:
    spline_data.append([x])
    
  for ys in newys:
    for i in range(len(ys)):
      spline_data[i].append(ys[i])
  
  return times, artist_stack, spline_data
    
    
def spline2_data(conn,type,user,interval=INTERVAL_RANGE,threshold=THRESHOLD,x_interval=0.1,stack=True,centered=True):
  times, artist_stack, data = raw_data(conn,type,user,interval,threshold,stack)
  if centered:
    center_data(data)

  import numpy, scipy

  x = numpy.array([d[0] for d in data])

  # Fit a spline to the original data and interpolate the curve with finer spacing
  spline_data = []
  newys = []
  newx = numpy.r_[0:data[-1][0]:x_interval]
  for i in range(len(artist_stack)+1):
    y = numpy.array([d[i+1] for d in data])
    
    import scipy.interpolate
    spline = scipy.interpolate.InterpolatedUnivariateSpline(x = x, y = y)
    newy = spline(newx)
    
    newys.append(newy)
    
  for x in newx:
    spline_data.append([x])
    
  for ys in newys:
    for i in range(len(ys)):
      spline_data[i].append(ys[i])
  
  return times, artist_stack, spline_data

#Returns the coordinates to place the artist label.  The point chosen is the point with the most plays in a given time period.  Also returns a recommended size to display the label at, based on the number of times that artist was played at the given point.
def label_positions_where_popular(data,artist_stack,text_size_func,min_size=7,max_size=30):
  results = []
  
  max_intervals = []
  for i in range(len(artist_stack)):
    max_interval = (0,0,0,0)
    for column in data[:-5]:
      if column[i+2] - column[i+1] > max_interval[0]:
        max_interval = (column[i+2] - column[i+1],column[0],column[i+1]+(column[i+2] - column[i+1])/2,i)
    max_intervals.append(max_interval)
  
  max_max = max(1,max([i[0] for i in max_intervals]))
  for max_interval in max_intervals:
    artist = artist_stack[max_interval[3]]
    size = max(min_size,max_size*max_interval[0]/max_max)
    
    width, height = text_size_func(artist)

    x, y = (max(0,max_interval[1]),max_interval[2])
    results.append((artist,x,y,size))
    
  return results

'''
def label_positions_where_popular(data,artist_stack,min_size=7,max_size=30):
  results = {}
  
  max_intervals = []
  for i in range(len(artist_stack)):
    max_interval = (0,0,0,0)
    for column in data[:-5]:
      if column[i+2] - column[i+1] > max_interval[0]:
        max_interval = (column[i+2] - column[i+1],column[0],column[i+1]+(column[i+2] - column[i+1])/2,i)
    max_intervals.append(max_interval)
  
  max_max = max(1,max([i[0] for i in max_intervals]))
  for max_interval in max_intervals:
    artist = artist_stack[max_interval[3]]
    size = max(min_size,max_size*max_interval[0]/max_max)

    x, y = (max(0,max_interval[1]),max_interval[2])
    results[artist] = (x,y,size)
    
  return results
'''

def label_positions_where_new(data,artist_stack,min_size=7,max_size=30):
  results = []
  
  max_intervals = []
  for i in range(len(artist_stack)):
    for column in data[:-5]:
      if column[i+2] - column[i+1] > 0:
        max_intervals.append((column[0],column[i+1],i))
        break
  
  for max_interval in max_intervals:
    artist = artist_stack[max_interval[2]]
    size = 9

    x, y = (max(0,max_interval[0]),max_interval[1])
    results.append((artist,x,y,size))
    
  return results
  

from collisions import Rect
def label_positions_no_collisions(data,artist_stack,text_size_func,min_size=7,max_size=30):
  print "Generating labels"

  results = []
  
  rects = []
 
  
  for i in range(len(artist_stack)):
    for column in data[:-5]:
      interval = column[i+2] - column[i+1]
      if interval > 5:
	artist = artist_stack[i]
	x, y = column[0],column[i+1]+(column[i+2] - column[i+1])/2
	width, height = text_size_func(artist)
	r = Rect(x,y,width,height)
	if not r.collides_with_set(rects):
	  results.append((artist,x,y,interval/2))
	  rects.append(r)

  return results